﻿using System;
using System.Linq;
using Algoritmia.DataStructures.Maps;
using Algoritmia.SemiRings;
using Algoritmia.StateSpaces;
using Algoritmia.Utils;

namespace Algoritmia.Schemes.DynamicProgramming
{
    public class IterativeDynamicProgrammingSolver<TState,TDecision,TResult,TWeight>
    {
        protected readonly Func<IBackwardsStateSpace<TState, TDecision>, IMap<TState, TResult>> _createMap;
        protected readonly Func<StateSpaceTopsorter<TState,TDecision>> _createStateSpaceTopsorter;

        public IterativeDynamicProgrammingSolver(Func<IBackwardsStateSpace<TState,TDecision>,IMap<TState,TResult>> createMap = null)
        {
            _createMap = createMap ?? (space => AlgoritmiaFactory.Maps<TState, TResult>.Default);
            _createStateSpaceTopsorter = () => new StateSpaceTopsorter<TState, TDecision>();
        }

        public TResult Solve(IBackwardsStateSpace<TState,TDecision> space,ISemiRing<TResult,TWeight> semiRing, Func<TState,TDecision,TWeight> wfn)
        {
            var mem = _createMap(space);
            var topsorter = _createStateSpaceTopsorter();
            foreach(var s in topsorter.Topsorted(space))
            {
                mem[s] = space.IsInitial(s) ? semiRing.One : semiRing.Zero;
                foreach(var d in space.IncomingDecisions(s))
                {
                    var s2 = space.Undo(s, d);
                    mem[s] = semiRing.Plus(mem[s], semiRing.Times(mem[s2], wfn(s2, d)));
                }
            }
            return semiRing.Sum(space.FinalStates.Select(f => mem[f]));
        }
    }

    public class IterativeDynamicProgrammingSolver<TState,TDecision,TResult> : 
        IterativeDynamicProgrammingSolver<TState,TDecision,TResult,TResult>
    {
        public IterativeDynamicProgrammingSolver(Func<IBackwardsStateSpace<TState, TDecision>, IMap<TState, TResult>> createMap = null)
            : base(createMap)
        {   
        }
    }
}
